import os
import subprocess

PREFIX = "arm-none-eabi-"
BUILDER = "DEV"

common_flags = []
build_projects = {}

build_projects["body"] = {
  "MAIN": "main.c",
  "STARTUP_FILE": "startup_stm32f413xx.s",
  "LINKER_SCRIPT": "stm32fx_flash.ld",
  "APP_START_ADDRESS": "0x08004000",
  "PROJECT_FLAGS": [
    "-mcpu=cortex-m4",
    "-mhard-float",
    "-DSTM32F4",
    "-DSTM32F413xx",
    "-mfpu=fpv4-sp-d16",
    "-fsingle-precision-constant",
    "-Os",
    "-g",
  ],
}

if os.getenv("RELEASE"):
  BUILD_TYPE = "RELEASE"
  cert_fn = os.getenv("CERT")
  assert cert_fn is not None, 'No certificate file specified. Please set CERT env variable'
  assert os.path.exists(cert_fn), 'Certificate file not found. Please specify absolute path'
else:
  BUILD_TYPE = "DEBUG"
  cert_fn = File("../certs/debug").srcnode().abspath
  common_flags += ["-DALLOW_DEBUG"]

if os.getenv("DEBUG"):
  common_flags += ["-DDEBUG"]

includes = [
  "inc/STM32F4xx_HAL_Driver/Inc",
  "inc",
  "..",
  ".",
]

c_sources = [
  ["hal_flash", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_flash.c"],
  ["hal_pwr", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_pwr.c"],
  ["hal_rcc", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_rcc.c"],
  ["hal_i2c", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_i2c.c"],
  ["hal_i2c_ex", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_i2c_ex.c"],
  ["hal_tim", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_tim.c"],
  ["hal_tim_ex", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_tim_ex.c"],
  ["hal_adc_ex", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_adc_ex.c"],
  ["hal_cortex", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_cortex.c"],
  ["hal_flash_ex", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_flash_ex.c"],
  ["hal_gpio", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_gpio.c"],
  ["hal_rcc_ex", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_rcc_ex.c"],
  ["hal", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal.c"],
  ["hal_adc", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_adc.c"],
  ["hal_dma", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_dma.c"],
  ["system", "inc/system_stm32f4xx.c"],
  ["it", "inc/stm32f4xx_it.c"],
  ["bldc", "bldc/bldc.c"],
  ["bldc_data", "bldc/BLDC_controller_data.c"],
  ["bldc_con", "bldc/BLDC_controller.c"],
  ["util", "util.c"],
]

c_bstub_sources = [
  ["hal_pwr", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_pwr.c"],
  ["hal_rcc", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_rcc.c"],
  ["hal_cortex", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_cortex.c"],
  ["hal_gpio", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_gpio.c"],
  ["hal_rcc_ex", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal_rcc_ex.c"],
  ["hal", "inc/STM32F4xx_HAL_Driver/Src/stm32f4xx_hal.c"],
  ["system", "inc/system_stm32f4xx.c"],
  ["it", "inc/stm32f4xx_it.c"],
  ["util", "util.c"],
]

def get_version(builder, build_type):
  try:
    git = subprocess.check_output(["git", "rev-parse", "--short=8", "HEAD"], encoding='utf8').strip()
  except subprocess.CalledProcessError:
    git = "00000000"
  return f"{git}"


def to_c_uint32(x):
  nums = []
  for _ in range(0x20):
    nums.append(x % (2**32))
    x //= (2**32)
  return "{" + 'U,'.join(map(str, nums)) + "U}"


def get_key_header(name):
  from Crypto.PublicKey import RSA

  public_fn = File(f'../certs/{name}.pub').srcnode().abspath
  rsa = RSA.importKey(open(public_fn).read())
  assert(rsa.size_in_bits() == 1024)

  rr = pow(2**1024, 2, rsa.n)
  n0inv = 2**32 - pow(rsa.n, -1, 2**32)

  r = [
    f"RSAPublicKey {name}_rsa_key = {{",
    f"  .len = 0x20,",
    f"  .n0inv = {n0inv}U,",
    f"  .n = {to_c_uint32(rsa.n)},",
    f"  .rr = {to_c_uint32(rr)},",
    f"  .exponent = {rsa.e},",
    f"}};",
  ]
  return r

def objcopy(source, target, env, for_signature):
    return '$OBJCOPY -O binary %s %s' % (source[0], target[0])

# Common autogenerated includes
git_ver = get_version(BUILDER, BUILD_TYPE)
with open("obj/gitversion.h", "w") as f:
  f.write(f'const uint8_t gitversion[8] = "{git_ver}";\n')

certs = [get_key_header(n) for n in ["debug", "release"]]
with open("obj/cert.h", "w") as f:
  for cert in certs:
    f.write("\n".join(cert) + "\n")

for project_name in build_projects:
  project = build_projects[project_name]
  linkerscript_fn = File(project["LINKER_SCRIPT"]).srcnode().abspath

  flags = [
    "-Wall",
    "-Wextra",
    "-Wstrict-prototypes",
    "-Werror",
    "-mlittle-endian",
    "-mthumb",
    "-nostdlib",
    "-fno-builtin",
    f"-T{linkerscript_fn}",
    "-std=gnu11",
    "-fdata-sections",
    "-ffunction-sections",
    "-Wl,--gc-sections",
  ] + project["PROJECT_FLAGS"] + common_flags

  project_env = Environment(
    ENV=os.environ,
    CC=PREFIX + 'gcc',
    AS=PREFIX + 'gcc',
    OBJCOPY=PREFIX + 'objcopy',
    OBJDUMP=PREFIX + 'objdump',
    ASCOM="$AS $ASFLAGS -o $TARGET -c $SOURCES",
    CFLAGS=flags,
    ASFLAGS=flags,
    LINKFLAGS=flags,
    LIBS = ['gcc',],
    CPPPATH=includes,
    BUILDERS={
      'Objcopy': Builder(generator=objcopy, suffix='.bin', src_suffix='.elf')
    }
  )
  startup = project_env.Object(project["STARTUP_FILE"])

  c_bstub_obj_list = []
  for c in c_bstub_sources:
    c_bstub_obj_list.append(project_env.Object(f"{c[0]}-{project_name}", c[1]))

  # Bootstub
  crypto_obj = [
    project_env.Object(f"rsa-{project_name}", "../crypto/rsa.c"),
    project_env.Object(f"sha-{project_name}", "../crypto/sha.c")
  ]
  bootstub_obj = project_env.Object(f"bootstub-{project_name}", "bootstub.c")
  bootstub_elf = project_env.Program(f"obj/bootstub.{project_name}.elf", [startup] + crypto_obj + c_bstub_obj_list + [bootstub_obj])
  bootstub_bin = project_env.Objcopy(f"obj/bootstub.{project_name}.bin", bootstub_elf)

  c_obj_list = []
  for c in c_sources:
    c_obj_list.append(project_env.Object(f"{c[0]}-{project_name}", c[1]))

  # Build main
  main_obj = project_env.Object(f"main-{project_name}", project["MAIN"])
  obj_list = [startup, main_obj] + c_obj_list
  main_elf = project_env.Program(f"obj/{project_name}.elf", obj_list,
    LINKFLAGS=[f"-Wl,--section-start,.isr_vector={project['APP_START_ADDRESS']}"] + flags)
  main_bin = project_env.Objcopy(f"obj/{project_name}.bin", main_elf)

  # Sign main
  sign_py = File("../crypto/sign.py").srcnode().abspath
  panda_bin_signed = project_env.Command(f"obj/{project_name}.bin.signed", main_bin, f"SETLEN=1 {sign_py} $SOURCE $TARGET {cert_fn}")
